
import os
import glob
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import (
    classification_report, accuracy_score, precision_score, recall_score, f1_score,
    mean_absolute_error
)

from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LinearRegression, LogisticRegression
from xgboost import XGBRegressor, XGBClassifier
from lightgbm import LGBMRegressor, LGBMClassifier
from catboost import CatBoostRegressor, CatBoostClassifier

# —— 特征组合定义 —— 
feature_sets = {
    'core5':   ['sdnn', 'rmssd', 'hf', 'lf', 'sampen'],
    'time2':   ['sdnn', 'rmssd'],
    'freq2':   ['hf', 'lf'],
    'mixed3':  ['rmssd', 'hf', 'sampen']
}

# —— 模型定义 ——
regressors = {
    'RandomForestReg': RandomForestRegressor(n_estimators=100, random_state=42),
    'LinearReg':       LinearRegression(),
    'XGBReg':          XGBRegressor(use_label_encoder=False, eval_metric='mae'),
    'LGBMReg':         LGBMRegressor(),
    'CatBoostReg':     CatBoostRegressor(verbose=0)
}
classifiers = {
    'RandomForestClf': RandomForestClassifier(n_estimators=100, random_state=42),
    'LogisticReg':     LogisticRegression(max_iter=1000),
    'XGBClf':          XGBClassifier(use_label_encoder=False, eval_metric='logloss'),
    'LGBMClf':         LGBMClassifier(),
    'CatBoostClf':     CatBoostClassifier(verbose=0)
}

# 构建特征矩阵
def build_features(feature_folder, min_windows=10, start_idx=2, end_idx=8):
    records = []
    for fp in glob.glob(os.path.join(feature_folder, '*_rrfeature.csv')):
        pid = os.path.basename(fp).replace('_rrfeature.csv','')
        df = pd.read_csv(fp)
        w = df[df['label']=='window']
        if len(w) < min_windows:
            continue
        sub = w.iloc[start_idx:end_idx]
        rec = {'patient_id': pid}
        for col in ['sdnn','rmssd','hf','lf','sampen']:
            vals = pd.to_numeric(sub[col], errors='coerce')
            rec[col] = vals.mean() if not vals.dropna().empty else np.nan
        records.append(rec)
    return pd.DataFrame(records)

# 读取年龄标签
def load_age(roster_csv):
    df = pd.read_csv(roster_csv, dtype=str)
    df['patient_id'] = df['userid'].str.strip() + '_' + df['recordid'].str.strip()
    df['age'] = pd.to_numeric(df['age'], errors='coerce')
    # 十年分段
    bins = list(range(0, 101, 10))
    labels = [f"{i}-{i+9}" for i in bins[:-1]]
    df['age_bin'] = pd.cut(df['age'], bins=bins, labels=labels, right=False)
    df = df.dropna(subset=['age_bin'])
    return df[['patient_id','age','age_bin']]

# 主程序
if __name__ == '__main__':
    feat_folder = '/database/home/duansizhang/hrv_predict/data/rr_featrue/'
    roster_csv  = '/database/private/mgcdb/info_xml_raw_match_1.csv'

    # 构建数据
    feat_df = build_features(feat_folder)
    age_df  = load_age(roster_csv)
    data    = feat_df.merge(age_df, on='patient_id', how='inner').dropna()
    print(f"共 {len(data)} 个合规样本")

    # 划分回归和分类的特征与标签
    X_base = data[['sdnn','rmssd','hf','lf','sampen']]
    y_reg   = data['age']
    y_clf   = LabelEncoder().fit_transform(data['age_bin'])

    # 回归评估
    print("\n===== 回归 MAE 评估 =====")
    for fs_name, feats in feature_sets.items():
        X = data[feats].replace([np.inf,-np.inf],np.nan).dropna()
        y = y_reg.loc[X.index]
        if X.empty: continue
        Xtr,Xte,ytr,yte = train_test_split(X, y, test_size=0.3, random_state=42)
        print(f"--- 特征组 {fs_name} ---")
        for name, reg in regressors.items():
            reg.fit(Xtr, ytr)
            yp = reg.predict(Xte)
            print(f"{name:14s} MAE {mean_absolute_error(yte,yp):.2f}")

    # 分类评估
    print("\n===== 年龄分段分类报告 =====")
    for fs_name, feats in feature_sets.items():
        X = data[feats].replace([np.inf,-np.inf],np.nan).dropna()
        y = y_clf[X.index]
        if X.empty: continue
        Xtr,Xte,ytr,yte = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
        print(f"--- 特征组 {fs_name} ---")
        for name, clf in classifiers.items():
            clf.fit(Xtr, ytr)
            yp = clf.predict(Xte)
            print(f"模型: {name}")
            print(classification_report(yte, yp, zero_division=0))

